import pandas as pd
import os

def read_and_clean(input_path: str) -> pd.DataFrame:
    """读取并清洗单个Excel文件"""
    # 读取第四个工作表（零部件汇总表）
    df = pd.read_excel(
        input_path,
        header=None,
    )

    # 反转数据行
    df = df.iloc[::-1].reset_index(drop=True)

    # 设置正确列名
    new_columns = df.iloc[0].tolist()
    df.columns = new_columns
    # 删除原始标题
    df = df[1:]
    return clean_dataframe(df)

def clean_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """数据清洗逻辑"""
    # 重命名列
    df.columns = [
        '序号', '名称及规格', '单位', '数量', 
        '材料', '单重', '总重', '备注'
    ]
    
    # 拆分名称及规格
    df[['名称', '型号及规范']] = df['名称及规格'].str.split(
        r'\s+', n=1, expand=True
    ).fillna('/')

    # 处理显示格式
    df['材料'] = df['材料'].apply(lambda x: x if pd.notnull(x) else '/')
    df['数量'] = df['数量'].apply(lambda x: x if pd.notnull(x) else '/')
    df['单位'] = df['单位'].apply(lambda x: x if pd.notnull(x) else '/')
    df['单重'] = df['单重'].apply(lambda x: x if pd.notnull(x) else '/')
    df['总重'] = df['总重'].apply(lambda x: x if pd.notnull(x) else '/')
    df['型号及规范'] = df['型号及规范'].apply(lambda x: x if pd.notnull(x) else '/')

    # 移除临时列
    df.drop(columns=['序号', '名称及规格'], inplace=True)

    return df

def merge_materials(df: pd.DataFrame) -> pd.DataFrame:
    """合并相同物料（处理整个目录数据）"""
    # 强制转换类型
    df['单重'] = pd.to_numeric(df['单重'], errors='coerce')
    df['数量'] = pd.to_numeric(df['数量'], errors='coerce')

    group_cols = ['名称', '型号及规范', '材料', '单重']
    agg_rules = {
        '数量': 'sum',
        '总重': 'sum',
        '备注': lambda x: ';'.join(x.dropna().unique()),
        '单位': 'first',
        '单重': 'first'
    }
    
    merged = df.groupby(group_cols, as_index=False).agg(agg_rules)

    # 重新计算总重(双校验机制)
    merged['总重'] = merged.apply(
        lambda row: row['单重'] * row['数量'] 
        if (pd.notnull(row['单重']) and pd.notnull(row['数量'])) 
        else row['总重'],  # 保留原始总重
        axis=1
    )

    return merged

def merge_number(df: pd.DataFrame) -> pd.DataFrame:
    """处理单重为空, 数量不为空, 材料名称相同的, 把他们合并"""
    # 强制转换类型
    df['数量'] = pd.to_numeric(df['数量'], errors='coerce')

    group_cols = ['名称', '型号及规范', '材料', '单位', '单重', '总重']
    agg_rules = {
        '数量': 'sum',
        '备注': lambda x: ';'.join(x.dropna().unique()),
    }

    merged = df.groupby(group_cols, as_index=False).agg(agg_rules)

    return merged

def format_columns(df: pd.DataFrame) -> pd.DataFrame:
    """格式化输出列"""
    column_order = [
        '名称', '型号及规范', '材料', '单位', 
        '数量', '单重', '总重', '备注'
    ]
    
    return df.reindex(columns=column_order)

def save_to_excel(df: pd.DataFrame, output_path: str) -> None:
    """保存Excel文件"""
    with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
        df.to_excel(writer, index=False)

def process_all_data(combined_df: pd.DataFrame) -> pd.DataFrame:
    """处理合并后的完整数据集"""
    # 清理材料 名称 为空的行
    combined_df = combined_df[combined_df['名称'] != '/']

    # 分离组合件
    combo_mask = (
        combined_df['单重'].isna() | 
        (combined_df['单重'] == '/') | 
        (combined_df['单重'] == '-') | 
        (combined_df['单重'] == '')
    )
    combo_df = combined_df[combo_mask].copy()
    normal_df = combined_df[~combo_mask].copy()
    
    # 合并普通材料
    merged_df = merge_materials(normal_df)
    # 合并单重为空的材料
    number_df = merge_number(combo_df)
    
    # 合并最终结果
    final_df = pd.concat([merged_df, number_df], ignore_index=True)
    return format_columns(final_df)

def process_directory(input_dir: str, output_path: str) -> None:
    """处理目录中的所有Excel文件"""
    all_dfs = []
    
    # 遍历目录收集所有数据
    for file_name in os.listdir(input_dir):
        if file_name.endswith('.xlsx') and not file_name.startswith('~$'):
            file_path = os.path.join(input_dir, file_name)
            print(f"正在处理: {file_name}")
            try:
                df = read_and_clean(file_path)
                all_dfs.append(df)
            except Exception as e:
                print(f"  错误：{str(e)}")
    
    if not all_dfs:
        print("未找到有效Excel文件")
        return
    
    # 合并所有数据
    combined_df = pd.concat(all_dfs, ignore_index=True)
    
    # 统一处理合并后的数据
    final_df = process_all_data(combined_df)
    
    # 保存结果
    save_to_excel(final_df, output_path)
    print(f"\n处理完成! 合并结果已保存至: {output_path}")

if __name__ == '__main__':
    input_directory = r'D:\Users\lijinsen\Desktop\dwg_excel'
    output_file = r'output\combined_output_dwg.xlsx'
    process_directory(input_directory, output_file)